// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief	Main Body of the PEPI_QT MPI Program
*/
/*! 
\defgroup	pepi_qt PEPI_QT - Epistasis for Quantitative Trait
\ingroup    mpi_stats
\brief		PEPI_QT - Epistasis for Quantitative Trait
\deprecated Program replaced by \ref pepi_qt_v2 which combines the jobs of statistical calculation and result merging.
\details
<p>
  This program is based on the epistasis (BT) %option of PLINK.
  It performs the same interaction calculation and unlike PLINK can also permute interaction calculations by distributing
  randomised data-sets across the cores of a computer cluster.
</p>

<p>This program is used in conjunction with \ref pepi_qt_merge and performs a permuted epistasis calculation which 
calculates the False Discovery Rates (FDR).</p>

<p>By setting an epistasis p-value cut-off threshold, this application will track interactions of interest in the permuted calculation.
Only interactions below the cut-off threshold are written to the results file.
The program has been tested, calculating the interactions of set up to 1000 SNPs.</p>

<h3>Usage</h3>
<p>The sequence of calls between the calculation and merge programs of epistasis (BT) are as follows:-</p>
<ul>
<li><strong>Call 1:-</strong><br>
pepi_qt.exe performs a non-permuted interaction calculations in a data-set
which is spread across N-cores.</li>
<li><strong>Call 2:-</strong><br>
pepi_qt_merge.exe (single-core) combines the interaction data from all the cores into a single
data file and identifies any interactions of interest based on a p-value significance threshold.</li>
<li><strong>Call 3:-</strong><br>
pepi_qt.exe runs permutations on a randomised data-sets, tracking interactions of interest 
and writing permutation counts to file. The number of required permutations is spread over N-cores.
</li>
<li><strong>Call 4:-</strong><br>
pepi_qt_merge.exe reads the permutation data from interesting interactions, 
combines this data and generates the final output file.</li>
</ul>

<p>If successful PEPI_QT returns with an exitcode of 0, if it encountered an error the programs returns with an exit code of -1.
This program does not create any output files and all temporary data files are deleted once all randomised calculations are complete.
</p>

<p>This documentation is vague. To get a handle on how to use this code, look at the relevant code snippet in the \ref examples
application which outlines the invocation of this epistasis program.</p>

<h2>OPTIONS</H2>

<table class=options border=0>
<tr><th class=param_header colspan=1>PARAMETER</th>
<th class=param_header colspan=3>DESCRIPTION</th></tr>

<tr><td class=param_name colspan=1>--id STRING</td>
<td class=param_desc colspan=3>Unique identifer for a cluster job (also sets output file stem).</td></tr>

<tr><td class=param_name colspan=1>--rdir PATH</td>
<td class=param_desc colspan=3>Path to remote working directory on central file system.</td></tr>

<tr><td class=param_name colspan=1>--ldir PATH</td>
<td class=param_desc colspan=3>Path to local working directory on each core of the cluster.</td></tr>

<tr><td class=param_name colspan=1>--bmap FILE</td>
<td class=param_desc colspan=3>Path to the BMAP file (binary-encoded SNP list).</td></tr>

<tr><td class=param_name colspan=1>--samples FILE</td>
<td class=param_desc colspan=3>Path to the Sample file (people and phenotype).</td></tr>

<tr><td class=param_name colspan=1>--iped FILE</td>
<td class=param_desc colspan=3>Path to the IPED file (Allelic data, Individual Major, bit-strings).</td></tr>

<tr><td class=param_name colspan=1>--nperm 1..10000</td>
<td class=param_desc colspan=3>
Specify the number of permutations to run on the cluster.
Specifying permutation number causes this programs to run randomised data-sets 
on each core of the cluster.</td></tr>

<tr><td class=param_name colspan=1>--missed-rank 0..101</td>
<td class=param_desc colspan=3>%Set processor rank to number specified by this parameter instead of rank assigned by MPI.<br>
To be used only if a section of a calculation needs to run a single core if a processor failed.
</td></tr>

<tr><td class=param_name colspan=1>--missed-procs 1..102</td>
<td class=param_desc colspan=3>%Set the No. processors associated with a calculation incase processor rank set by <em>missed-rank</em>.</td></tr>

</table>
*/

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "bmapfile.h"
#include "chromosomes.h"
#include "copyfile.h"
#include "doublefile.h"
#include "getopt.h"
#include "ipedfile.h"
#include "linear.h"
#include "mpi_error.h"
#include "mt19937ar.h"
#include "options.h"
#include "permfile.h"
#include "pmaths.h"
#include "sample.h"
#include "seedfile.h"
#include "suffixes.h"
#include "utils.h"

//#define ECHO_PERMUTATION

/*!
\brief Command-line options of PEPI_QT
*/
struct option long_options[] = {
	{ "bmap", required_argument, NULL, 'm' }, // Path to the map file.
	{ "samples", required_argument, NULL, 's' }, // Path to the map file.
	{ "iped", required_argument, NULL, 'p' }, // Path to the PED file.
	{ "id", required_argument, NULL, 'i' }, // Run ID
	{ "rdir", required_argument, NULL, 'r'}, // Remote Output direct
	{ "ldir", required_argument, NULL, 'l'}, // Local Output direct
	{ "sex-effect", optional_argument, NULL, 'x'}, // Local Output direct
	{ "missed-rank", required_argument, NULL, 'n'}, // A rank that failed for whatever reason that needs re-running.
	{ "missed-procs", required_argument, NULL, 'o'}, // No. processors associated with the original calculation.
	{ "nperm", required_argument, NULL, 'P'}, // No. permutations.
	{ "epi_alpha1", required_argument, NULL, 'e' }, // Tight EPI threshold
	{ "epi_alpha2", required_argument, NULL, 'f' }, // Tight EPI threshold
	{ 0,    0,    0,    0   } 
};

/*!
\brief Main body of the PEPI_QT program
\ingroup pepi_qt
\deprecated Program replaced by \ref pepi_qt_v2 which combines the jobs of statistical calculation and result merging.
*/
int pepi_qt_main(int argc, char *argv[]) {
	char ch, *ipedData = NULL;
	struct selected_options ops;
	int myrank, numprocs, result;
	struct chromosomes chromos;
	int nRequiredArgs = 6;
	struct small_locus *map = NULL;
	struct sample *samples = NULL;
	int nMap = 0, nSamples = 0;
	int ipedNSamples = 0, nSNPs = 0, nIpedRecordSizeInBytes = 0;
	BOOL jobDone = FALSE;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];
	int indices[2], missedRank = -1, missedProcs = -1;
	int *summary_sig = NULL, *summary_good = NULL, *best_partner = NULL, e1, nperm = -1;
	register int e2;
	double *best_score = NULL; 
	int *seeds = NULL, nSeeds = 0, *perms = NULL, *best_partner_permuted_counts = NULL,  nOriginal_best_partner_scores = 0;
	double *original_best_partner_scores = NULL; // Best-pairwise match pvalues from the original data-set.
	struct INT2 *interesting_interactions = NULL;
	int nInterestingInteractions = 0;

	// Initialise MPI environment, get size and rank.
	if ((result = MPI_Init(&argc,&argv)) != MPI_SUCCESS) 
		mpi_error(result, "Initialisation");
	if ((result = MPI_Comm_size(MPI_COMM_WORLD,&numprocs)) != MPI_SUCCESS)
		mpi_error(result, "Number of processors");
	if ((result = MPI_Comm_rank(MPI_COMM_WORLD,&myrank)) != MPI_SUCCESS)
		mpi_error(result, "Processor number");

	options_init(&ops);
	chromosomes_init(&chromos);
	chromosomes_define_human(&chromos);

	// Want to set these set via the command-line.
	ops.ldir[0] = ops.szRunId[0] = '\0'; 
	ops.verbose = FALSE;

	while ((ch = (char) getopt_long(argc, argv, "m:v:p:f:i:", long_options, NULL)) != -1) {
		double value;

		switch (ch) {
			case 'i':
				strcpy(ops.szRunId, optarg);
				nRequiredArgs--;
				break;
			case 's': // Sample File
				strcpy(ops.szSampleFileName, optarg);
				nRequiredArgs--;
				break;
			case 'p': // IPED file
				strcpy(ops.szIpedFileName, optarg);
				nRequiredArgs--;
				break;
			case 'm': // Bmap File
				strcpy(ops.szBmapFileName, optarg);
				nRequiredArgs--;
				break;
			case 'r': // Remote output directory
				strcpy(ops.rdir, optarg);
				nRequiredArgs--;
				break;
			case 'l': // Local output directory
				strcpy(ops.ldir, optarg);
				nRequiredArgs--;
				break;
			case 'x': // Sex Effect
				ops.sex_effect = TRUE;
				break;
			case 'n': // Missed rank to re-run calculation on a single processor.
				missedRank = atoi(optarg);
				break;
			case 'o': // No. processors associated with a 'missed' rank.
				missedProcs = atoi(optarg);
				break;
			case 'P': // Number of Permutations.
				nperm = atoi(optarg);
				if (nperm == 0) nperm = 10000;
				break;
			case 'e': // Tight EPI threshold (epi_alpha1).
				value = atof(optarg);
				if (!(value > 0 && value <= 1.0)) value = DEFAULT_EPI_ALPHA1;
				ops.epi_alpha1 = value;
				break;
			case 'f': // Wide EPI threshold (epi_alpha2).
				value = atof(optarg);
				if (!(value > 0 && value <= 1.0)) value = DEFAULT_EPI_ALPHA2;
				ops.epi_alpha2 = value;
				break;
			default:
				break;
		}
	}

	// If the required arguements are not specified, there's no point carrying on.
	if (nRequiredArgs != 0) {
		chromosomes_free_attributes(&chromos);
		MPI_Finalize();
		return -1;
	}

	// See if a calculation needs to be re-run on a 'single' processor.
	// Check that all the command-line options are set.
	if (missedRank >= 0 || missedProcs >= 0) {
		BOOL exitEarly = FALSE;

		if (numprocs != 1) exitEarly = TRUE;
		else if (missedRank == -1 || missedProcs == -1) exitEarly = TRUE;
		else if (!(missedRank >= 0 && missedRank <= missedProcs)) exitEarly = TRUE;

		if (exitEarly) {
			chromosomes_free_attributes(&chromos);
			MPI_Finalize();
			return -1;
		} else {
			myrank = missedRank;
			numprocs = missedProcs;
		}
	}

	// The non-permutated calculation.
	if (nperm == -1) {
		if ((map = bmapfile_load(&ops, &nMap)) != NULL && 
			(samples = samplefile_load(&ops, &nSamples)) != NULL &&
			(ipedData = ipedfile_load(&ops, &ipedNSamples, &nSNPs, &nIpedRecordSizeInBytes)) != NULL) 
		{
			ops.epi_alpha1 = fabs(ltqnorm(ops.epi_alpha1 / 2));
			ops.epi_alpha2 = fabs(ltqnorm(ops.epi_alpha2 / 2));

			// The result arrays for the best pair-wise interactions found on this rank.
			summary_sig = create_int_array(nSNPs, 0);
			summary_good = create_int_array(nSNPs, 0);
			best_partner = create_int_array(nSNPs, -1);
			best_score = create_vector(nSNPs, 0);

			// Magic numbers match so nothing weird is going on.
			if (nMap == nSNPs && ipedNSamples == nSamples && getDHLoopLimits(nSNPs, numprocs, myrank, indices) &&
				summary_sig != NULL && summary_good != NULL && best_partner != NULL && best_score != NULL) 
			{
				struct linear_model lm;
				FILE *out = NULL;

				// Create the local/remote output file paths
				set_datafile_path(localFilepath, ops.ldir, ops.szRunId, myrank, SUFFIX_LOCAL_EPI_FILE);
				set_datafile_path(remoteFilepath, ops.rdir, ops.szRunId, myrank, SUFFIX_LOCAL_EPI_FILE);

				// Initialise the model & create the local output file.
				out = fopen(localFilepath, "w");
				linear_model_init(&lm);
				linear_model_set_np(&lm, &ops);
				if (linear_model_alloc(&lm, nSamples) && out != NULL) {
					// Apply original order to the phenotype array.
					samplefile_set_original_order(samples, nSamples);

					fprintf(out, "%i\n", numprocs); // Add a bit of header info to the output file.

					// Do the pair-wise epistasis things.
					for (e1 = indices[0]; e1 < indices[1]; e1++) {
						for (e2 = 0; e2 < nMap; e2++) {
							double chisq, z;
							int testParameter = 3; // interaction
							
							if (e1 >= e2) continue;
							else if (e1==e2) continue;

							// Analyse diploid autosomal SNPs
							if (!(chromos.chr_sex[map[e1].chr] || chromos.chr_sex[map[e2].chr] || 
								chromos.chr_haploid[map[e1].chr] || chromos.chr_haploid[map[e2].chr])) 
							{
								// Build design matrix
								samplefile_set_missing(lm.miss, lm.nMiss, samples);

								// Main effect of SNP 1
								linear_model_add_additive_snp(&lm, &chromos, map, e1, 0);

								// Main effect of SNP 2
								linear_model_add_additive_snp(&lm, &chromos, map, e2, 1);

								// Epistasis (use a single interaction object).
								lm.nInteraction = 1;
								lm.interaction[0].p1 = 1;
								lm.interaction[0].p2 = 2;

								linear_model_build_design_matrix(&lm, ipedData, samples, nSamples, nIpedRecordSizeInBytes, &ops);								
								linear_model_fitLM(&lm);
								linear_model_validParameters(&lm);
								chisq = linear_model_get_statistic(&lm, testParameter);
								z = sqrt(chisq);

								// Is this result worth displaying?
								if (lm.all_valid) {
									// Count as a good result
									summary_good[e1]++;
									summary_good[e2]++;

									// Do we want to record this as part of the summary for the first set? 	  
									if (z >= ops.epi_alpha2) {
										summary_sig[e1]++;
										summary_sig[e2]++;
									}

									// Is this result the best scrore yet for marker in set A?
									if (z > best_score[e1]) {
										best_score[e1] = z;
										best_partner[e1] = e2;
									}

									if (z > best_score[e2]) {
										best_score[e2] = z;
										best_partner[e2] = e1;
									}
								}

								// Is this worth recording?
								if (z >= ops.epi_alpha1 && lm.all_valid) fprintf(out, "%i,%i,%10.16f,%10.16f\n", e1, e2, lm.coef[3], z);
							}
						}
					}
					if (out != NULL) { fclose(out); out = NULL; }
					if (copyAndDeleteFile(localFilepath, remoteFilepath)) jobDone = TRUE; 
					
					// Write out the summary output file of the best interaction for each SNP.
					if (jobDone) {
						jobDone = FALSE;

						// Create the local/remote output file paths
						set_datafile_path(localFilepath, ops.ldir, ops.szRunId, myrank, SUFFIX_LOCAL_BEPI_FILE);
						set_datafile_path(remoteFilepath, ops.rdir, ops.szRunId, myrank, SUFFIX_LOCAL_BEPI_FILE);
						out = fopen(localFilepath, "w");
						if (out != NULL) {
							for (e1 = 0; e1 < nSNPs; e1++) {
								fprintf(out, "%i,%i,%i,", e1, summary_sig[e1], summary_good[e1]);

								if (realnum(best_score[e1])) fprintf(out, "%f,", best_score[e1]);
								else fprintf(out, "NA,");

								if (best_partner[e1] != -1) 
									fprintf(out,"%i\n",best_partner[e1]); 
								else
									fprintf(out,"NA,\n");

							}
							fclose(out);
							if (copyAndDeleteFile(localFilepath, remoteFilepath)) jobDone = TRUE; 
						}
					}
				}
				linear_model_free_attributes(&lm); 
			}
		}
	} else {
		// Do Epistasis on permuted data-sets.
		if ((map = bmapfile_load(&ops, &nMap)) != NULL && 
			(samples = samplefile_load(&ops, &nSamples)) != NULL &&
			(ipedData = ipedfile_load(&ops, &ipedNSamples, &nSNPs, &nIpedRecordSizeInBytes)) != NULL &&
			(seeds = seedfile_load(&ops, &nSeeds)) != NULL)  
		{
			int nPermsPerProcessor = get_tasks_per_processor(numprocs, nperm);
			int start_perm, end_perm, i;

			// Calculate indexes for the outer permutation loop.
			start_perm = myrank * nPermsPerProcessor;
			end_perm = start_perm + nPermsPerProcessor;
			if (end_perm > nperm) end_perm = nperm;

			// Epistasis with permutation, N/B must limit the input size.
			best_score = create_vector(nSNPs, 0);
			perms = (int*) malloc(nSamples * sizeof(int)); // Permutation indices.
			best_partner_permuted_counts = create_int_array(nSNPs, 0);

			// Load the best partner data from the ordinary calculation.
			original_best_partner_scores = doublefile_read(&ops, SUFFIX_DOUBLE_FILE, &nOriginal_best_partner_scores);

			// Load the Interaction file containing interesting co-ordinates.
			sprintf(remoteFilepath, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_EPI_FILE2);
			interesting_interactions = int2_load(remoteFilepath, &nInterestingInteractions);
			if (interesting_interactions != NULL && nInterestingInteractions > 0) qsort(interesting_interactions, nInterestingInteractions, sizeof(struct INT2), int2_cmp);

			// Magic numbers match so nothing weird is going on.
			if (nMap == nSNPs && ipedNSamples == nSamples && nSeeds >= numprocs && perms != NULL &&
				best_partner_permuted_counts != NULL && original_best_partner_scores != NULL && best_score != NULL) 
			{
				struct linear_model lm;
				FILE *out = NULL;

				samplefile_set_original_order(samples, nSamples);
				
				// Create the local/remote output file paths
				set_datafile_path(localFilepath, ops.ldir, ops.szRunId, myrank, SUFFIX_RBEPI_FILE);
				set_datafile_path(remoteFilepath, ops.rdir, ops.szRunId, myrank, SUFFIX_RBEPI_FILE);

				// Permutation Loop
				linear_model_init(&lm);
				linear_model_set_np(&lm, &ops);

				if (linear_model_alloc(&lm, nSamples)) {
					init_genrand(seeds[myrank]);
					for (i = start_perm; i < end_perm; i++) {
						double zMaxForWholeDataSet = -1000;

#ifdef ECHO_PERMUTATION
					if (myrank == ROOT) {
						fprintf(stdout, "Permutation: %i\r", i + 1);
						fflush(stdout);
					}
#endif
						// Randomise the sample phenotype.
						permute(perms, nSamples);
						samplefile_permute_phenotype(samples, perms, nSamples);

						// Do epistasis on the randomised data-set.
						zero_vector(best_score, nSNPs);
						int2_reset_valid(interesting_interactions, nInterestingInteractions);

						for (e1 = 0; e1 < nMap; e1++) {
							for (e2 = 0; e2 < nMap; e2++) {
								double chisq, z;
								int testParameter = 3; // interaction
								
								if (e1 >= e2) continue;
								else if (e1==e2) continue;

								// Analyse diploid autosomal SNPs
								if (!(chromos.chr_sex[map[e1].chr] || chromos.chr_sex[map[e2].chr] || 
									chromos.chr_haploid[map[e1].chr] || chromos.chr_haploid[map[e2].chr])) 
								{
									// Build design matrix
									samplefile_set_missing(lm.miss, lm.nMiss, samples);

									// Main effect of SNP 1
									linear_model_add_additive_snp(&lm, &chromos, map, e1, 0);

									// Main effect of SNP 2
									linear_model_add_additive_snp(&lm, &chromos, map, e2, 1);

									// Epistasis (use a single interaction object).
									lm.nInteraction = 1;
									lm.interaction[0].p1 = 1;
									lm.interaction[0].p2 = 2;

									linear_model_build_design_matrix(&lm, ipedData, samples, nSamples, nIpedRecordSizeInBytes, &ops);								
									linear_model_fitLM(&lm);
									linear_model_validParameters(&lm);
									chisq = linear_model_get_statistic(&lm, testParameter);
									z = sqrt(chisq);

									if (lm.all_valid) {
										if (z > best_score[e1]) best_score[e1] = z;
										if (z > best_score[e2]) best_score[e2] = z;
										if (z > zMaxForWholeDataSet) zMaxForWholeDataSet = z;

										// Update the 'interesting interactions' result array.
										if (z >= ops.epi_alpha && interesting_interactions != NULL) {
											int index = -1;
											struct INT2 interaction;

											interaction.p1 = e1;
											interaction.p2 = e2;
											index = int2_search(interesting_interactions, nInterestingInteractions, &interaction); 
											if (index >= 0) {
												interesting_interactions[index].valid = TRUE;
												interesting_interactions[index].nPermutations++;
												if (z > interesting_interactions[index].z) interesting_interactions[index].R++; // EMP1 counts
											}
										}
									}
								}
							}
						}

						// Update the best_score permutation counts.
						for (e1 = 0; e1 < nSNPs; e1++) {
							if (best_score[e1] > 0) {
								if (best_score[e1] > original_best_partner_scores[e1]) best_partner_permuted_counts[e1]++;
							}
						}

						// EMP2 counts
						if (interesting_interactions) {
							for (e1 = 0; e1 < nInterestingInteractions; e1++) {
								if (interesting_interactions[e1].valid) {
									if (zMaxForWholeDataSet >= interesting_interactions[e1].z) interesting_interactions[e1].maxR++;
								}
							}
						}
					}
					linear_model_free_attributes(&lm);

					// Create the local output file.
					set_datafile_path(localFilepath, ops.ldir, ops.szRunId, myrank, SUFFIX_RBEPI_FILE);
					set_datafile_path(remoteFilepath, ops.rdir, ops.szRunId, myrank, SUFFIX_RBEPI_FILE);
					out = fopen(localFilepath, "w");
					if (out != NULL) {
						// Header lines, numprocs and the number of permutations this rank is running.
						fprintf(out, "%i\n", numprocs);
						fprintf(out, "%i\n", abs(end_perm - start_perm));
						for (e1 = 0; e1 < nSNPs; e1++) fprintf(out, "%i\n", best_partner_permuted_counts[e1]);
						fclose(out);
						if (copyAndDeleteFile(localFilepath,remoteFilepath)) jobDone = TRUE;
					} 

					// Create the local output file for the randomised significant hits (but only if there's data).
					if (interesting_interactions && jobDone) {
						jobDone = FALSE;
						set_datafile_path(localFilepath, ops.ldir, ops.szRunId, myrank, SUFFIX_REPI_FILE);
						set_datafile_path(remoteFilepath, ops.rdir, ops.szRunId, myrank, SUFFIX_REPI_FILE);
						out = fopen(localFilepath, "w");
						if (out != NULL) {
							// Header lines, numprocs and the number of permutations this rank is running.
							fprintf(out, "%i\n", numprocs);
							fprintf(out, "%i\n", abs(end_perm - start_perm));
							for (e1 = 0; e1 < nInterestingInteractions; e1++) {		
								fprintf(out, "%i,%i,%i,%i,%i\n", interesting_interactions[e1].p1,interesting_interactions[e1].p2, interesting_interactions[e1].R, interesting_interactions[e1].maxR, interesting_interactions[e1].nPermutations);
							}
							fclose(out);
							if (copyAndDeleteFile(localFilepath,remoteFilepath)) jobDone = TRUE;
						}
					}
				}
			}
		}
	}

	// Free now as not needed anymore
	if (samples != NULL) free(samples);
	if (ipedData != NULL) free(ipedData);
	if (map != NULL) free(map); 
	if (summary_sig != NULL) free(summary_sig);
	if (summary_good != NULL) free(summary_good);
	if (best_partner != NULL) free(best_partner);
	if (best_score != NULL) free(best_score);
	if (seeds != NULL) free(seeds);
	if (perms != NULL) free(perms);
	if (original_best_partner_scores != NULL) free(original_best_partner_scores);
	if (best_partner_permuted_counts != NULL) free(best_partner_permuted_counts);
	if (interesting_interactions != NULL) free(interesting_interactions);
	chromosomes_free_attributes(&chromos);

	MPI_Finalize();

	if (jobDone) return 0;
	else return -1;
}

int main(int argc, char *argv[]) {
	return pepi_qt_main(argc, argv);
}